import traci

from src.sumo_simulation.sumo_manager import SumoSimulation as Sumo

from src.simulation_system.platoon import Platoon

from src.simulation_system.utils import compute_distance
from src.simulation_system.vehicle import Vehicle


class CavWorld(object):
    """
    A customized world object to save all CDA vehicle
    information and shared ML models. During co-simulation,
    it is also used to save the sumo-carla id mapping.

    Attributes
    ----------
    _vehicle_id_set : set
        A set that stores vehicle IDs.

    _single_vehicle_id_set : set
        A set that stores single_vehicle IDs.

    _vehicle_manager_dict : dict
        A dictionary that stores vehicle managers.

    _platooning_manager_dict : dict
        A dictionary that stores platooning managers.

    """

    def __init__(self, platoon):
        self.platoon = platoon
        self._vehicle_id_set = set()
        self._single_vehicle_id_set = set()
        self._vehicle_manager_dict = {}
        self._traffic_lights = {}
        if platoon:
            self._platooning_manager_dict = {}

    def add_trafficlight(self, traffic_light):
        tl_id = traffic_light.get_tl_id()
        self._traffic_lights[tl_id] = traffic_light

    def get_trafficlight(self):
        return self._traffic_lights

    def get_trafficlight_data(self):
        tl_data = {}
        for tl_id, tl in self._traffic_lights.items():
            tl_data[tl.get_tl_id()] = tl.get_states()
        return tl_data

    def update_trafficlight(self):
        for tl_id, traffic_light in self._traffic_lights.items():
            traffic_light.update_info()

    def add_vehicle(self, vehicle_id):
        vm = Vehicle(vehicle_id, self, self.platoon)
        self._vehicle_id_set.add(vehicle_id)
        self._single_vehicle_id_set.add(vehicle_id)
        self._vehicle_manager_dict[vehicle_id] = vm

    def add_platoon(self, start_vehicle):
        pm = Platoon([start_vehicle], self)
        vid = start_vehicle.get_vehicle_id()
        if vid in self._single_vehicle_id_set:
            self._single_vehicle_id_set.remove(vid)
        self._platooning_manager_dict[pm.get_platoon_id()] = pm

    def destroy_vehicle(self, vehicle_id):
        vm = self._vehicle_manager_dict[vehicle_id]
        if self.platoon:
            pid = vm.get_platoon_id()
            if pid is not None:
                pm = self._platooning_manager_dict[pid]
                if pm.get_platoon_len() == 1:
                    self._platooning_manager_dict.pop(pid)
                else:

                    vehicles = pm.get_vehicles()
                    vehicles.remove(vm)
                    new_pm = Platoon(vehicles, self)
                    self._platooning_manager_dict.pop(pid)
                    self._platooning_manager_dict[new_pm.get_platoon_id()] = new_pm
        if vehicle_id in self._vehicle_id_set:
            self._vehicle_id_set.remove(vehicle_id)
        if vehicle_id in self._single_vehicle_id_set:
            self._single_vehicle_id_set.remove(vehicle_id)
        if vehicle_id in self._vehicle_manager_dict.keys():
            self._vehicle_manager_dict.pop(vehicle_id)

    def destroy_platooning(self, pid):
        if pid in self._platooning_manager_dict.keys():
            vehicles = self.get_platooning_manager(pid).get_vehicles()
            self._platooning_manager_dict.pop(pid)

            for vehicle in vehicles:
                self._single_vehicle_id_set.add(vehicle.get_vehicle_id())

    def split_platoon(self, pid, index):
        members = self._platooning_manager_dict[pid].get_vehicles()
        member1 = members[0:index]
        member2 = members[index:]
        pm1 = Platoon(member1, self)
        pm2 = Platoon(member2, self)
        self._platooning_manager_dict[pid].set_active(False)
        self._platooning_manager_dict[pm1.get_platoon_id()] = pm1
        self._platooning_manager_dict[pm2.get_platoon_id()] = pm2
        pm1.set_merge()
        pm2.set_merge()

    def merge_platoon(self, pid1, pid2):
        member1 = self._platooning_manager_dict[pid1].get_vehicles()
        member2 = self._platooning_manager_dict[pid2].get_vehicles()
        if self._platooning_manager_dict[pid1].get_leader().get_lane_pos() > \
                self._platooning_manager_dict[pid2].get_leader().get_lane_pos():
            member1 = member1 + member2
            pm = Platoon(member1, self)
        else:
            member2 = member2 + member1
            pm = Platoon(member2, self)
        self._platooning_manager_dict[pid1].set_active(False)
        self._platooning_manager_dict[pid2].set_active(False)
        self._platooning_manager_dict[pm.get_platoon_id()] = pm
        pm.set_merge()

    def get_single_vehicle_ids(self):
        return list(self._single_vehicle_id_set)

    def get_vehicle_ids(self):
        return list(self._vehicle_id_set)

    def get_platoon_ids(self):
        return list(self._platooning_manager_dict.keys())

    def get_vehicle_manager(self, vehicle_id):
        if vehicle_id in self._vehicle_manager_dict.keys():
            return self._vehicle_manager_dict[vehicle_id]

    def get_platoon_data_by_vid(self, vehicle_id):
        if self.platoon:
            for pid, pm in self._platooning_manager_dict.items():
                if pm.if_in_platoon(vehicle_id):
                    return {"Color": pm.get_color(), "Members": pm.get_vehicle_ids()}
        return None

    def get_platooning_manager(self, pid):
        if pid in self._platooning_manager_dict.keys():
            return self._platooning_manager_dict[pid]

    def get_nearby_cav_data(self, vehicle_id, communication_range):
        nearby_cav = dict()
        current_vehicle = self.get_vehicle_manager(vehicle_id)
        for vid, vm in self._vehicle_manager_dict.items():
            # avoid add itself as the cav nearby
            if vid == vehicle_id:
                continue
            distance = compute_distance(
                current_vehicle.get_pos(),
                vm.get_pos())
            if distance < communication_range:
                # 获取当前遍历车辆的车道ID

                other_lane_id = vm.get_lane_id()
                # #判断是否在同一车道且同向
                # if if_same_lane(current_vehicle.get_lane_id(), other_lane_id):
                #     # 如果是，则添加到 nearby_cav 字典中
                #     pid = vm.get_platoon_id()
                #     nearby_cav[vid] = pid
                # 判断是否在同一车道
                other_road_id = vm.get_road_id()
                if if_same_road(current_vehicle.get_road_id(), other_road_id):
                    pid = vm.get_platoon_id()
                    nearby_cav[vid] = pid

        return nearby_cav

    def get_lane_vehicles(self, lane_id):
        vehicles = []
        for vid, vm in self._vehicle_manager_dict.items():
            if vm.get_lane_id() == lane_id:
                vehicles.append(vm)
        return vehicles

    def request_join(self, vehicle_id, pid):
        pm = self.get_platooning_manager(pid)
        vm = self.get_vehicle_manager(vehicle_id)
        if pm.if_join(vm):
            vehicles = pm.get_vehicles()
            vehicles.append(vm)
            self._platooning_manager_dict.pop(pid)
            self._single_vehicle_id_set.remove(vehicle_id)
            pm = Platoon(vehicles, self)
            pm.set_merge()
            self._platooning_manager_dict[pm.get_platoon_id()] = pm

    def get_vehicles_pos(self):
        vehicles_pos = {}
        for vid, vm in self._vehicle_manager_dict.items():
            vehicles_pos[vid] = {
                "pos": list(vm.get_pos()),  # 位置应该是一个数组，如 [float, float, float]
                "rot": vm.get_rotation()  # 旋转可以是一个单独的浮点数，或者如果需要的话，也可以是一个数组
            }
        return vehicles_pos

    def get_vehicles_data(self):
        vehicles_data = {}
        for vid, vm in self._vehicle_manager_dict.items():
            vehicles_data[vid] = {
                "pid": vm.get_platoon_id(),
                "pos": vm.get_pos(),  # 位置应该是一个数组，如 [float, float, float]
                "rot": vm.get_rotation()  # 旋转可以是一个单独的浮点数，或者如果需要的话，也可以是一个数组
            }
        return vehicles_data

    def get_lane_density(self):
        lane_density = {}
        for vid, vm in self._vehicle_manager_dict.items():
            if vm.get_lane_id() not in lane_density.keys():
                lane_density[vm.get_lane_id()] = 1
            else:
                lane_density[vm.get_lane_id()] += 1
        return lane_density

    def get_platoon_data(self):
        platoon_data = {}
        for pid in self._platooning_manager_dict.keys():
            vehicles = self.get_platooning_manager(pid).get_vehicles()
            vehicles_id = []
            for vehicle in vehicles:
                vehicles_id.append(vehicle.get_vehicle_id())
            platoon_data[pid] = vehicles_id
        return platoon_data

    def update_active_platoon(self):
        un_platoon = []
        for pid, pm in self._platooning_manager_dict.items():
            if not pm.get_active():
                un_platoon.append(pid)
        for pid in un_platoon:
            self._platooning_manager_dict.pop(pid)


def if_same_lane(lane_id1, lane_id2):
    """判断两辆车是否在同一车道且同向行驶。"""
    if lane_id1 != lane_id2:
        return False
    else:
        return True


def if_same_road(road_id1, road_id2):
    if road_id1 != road_id2:
        return False
    else:
        return True
